import asyncio
import logging
from mcp import ClientSession, stdio_client, StdioServerParameters

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("calculation_mcp_client")

class CalculationClient:
    def __init__(self):
        self.session = None
        self.running = False

    async def start(self):
        try:
            # 使用标准输入输出连接本地服务器
            server_params = StdioServerParameters(
                command="python",
                args=["CalculationServer.py"]
            )
            
            # 创建stdio上下文和会话
            stdio_context = stdio_client(server_params)
            read, write = await stdio_context.__aenter__()
            self.session = ClientSession(read, write)
            await self.session.__aenter__()
            
            # 初始化会话
            await self.session.initialize()
            self.running = True
            
            print("已连接到本地计算服务器！")
            print("可用工具: add, subtract, multiply, divide")
            print("输入命令（如 /add 2 3），输入 exit 退出")
            
            # 启动输入循环
            await self.input_loop()
        except Exception as e:
            logger.error(f"连接服务器失败: {e}")
            raise
        finally:
            await self.cleanup()
            print("已断开连接")

    async def call_tool(self, tool_name: str, *args):
        """调用服务器端的工具函数"""
        if not self.session or not self.running:
            raise RuntimeError("客户端未连接或已停止")
            
        # 构建工具调用参数
        if len(args) == 1:
            arguments = {"a": args[0]}
        elif len(args) >= 2:
            arguments = {"a": args[0], "b": args[1]}
        else:
            arguments = {}
            
        logger.info(f"调用工具: {tool_name}，参数: {arguments}")
        
        # 调用服务端工具
        result = await self.session.call_tool(
            name=tool_name,
            arguments=arguments
        )
        
        # 解析结果
        if hasattr(result, 'isError') and not result.isError:
            if hasattr(result, 'content') and result.content and len(result.content) > 0:
                if hasattr(result.content[0], 'text'):
                    return result.content[0].text
                return str(result.content[0])
            return str(result)
        else:
            error_msg = "工具调用失败"
            if hasattr(result, 'content') and result.content and len(result.content) > 0:
                error_msg += f": {str(result.content[0])}"
            raise Exception(error_msg)

    async def input_loop(self):
        """处理用户输入，调用对应的工具函数"""
        while self.running:
            cmd = await asyncio.to_thread(input, "输入命令: ")
            if cmd.lower() == "exit":
                self.running = False
                break
            
            # 解析命令（如 /add 2 3 -> 调用add工具，参数2和3）
            if cmd.startswith("/"):
                parts = cmd.split()
                if len(parts) < 2:
                    print("命令格式错误，请使用: /工具名 参数1 参数2...")
                    continue
                
                tool_name = parts[0][1:]  # 去除斜杠
                try:
                    # 转换参数为数字
                    args = [float(x) for x in parts[1:]]
                    
                    # 调用服务端的工具函数
                    result = await self.call_tool(tool_name, *args)
                    print(f"计算结果: {result}")
                except ValueError as e:
                    print(f"参数错误: {e}")
                except Exception as e:
                    print(f"调用失败: {e}")
            else:
                print(f"忽略非工具命令: {cmd}")

    async def cleanup(self):
        """清理资源"""
        self.running = False
        
        # 关闭会话
        if self.session:
            try:
                await self.session.__aexit__(None, None, None)
            except Exception as e:
                logger.warning(f"关闭会话时出错: {e}")
            finally:
                self.session = None

# 主函数
async def main():
    # 连接到本地服务器
    client = CalculationClient()
    await client.start()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("客户端已停止")